#!/usr/bin/env python3
"""
MCP端LLM+RAG+Langraph集成处理逻辑
处理从LLM系统接收的分析结果，智能选择子Agent并执行决策
"""

import asyncio
import aiohttp
import json
from datetime import datetime
from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from enum import Enum

class IssueType(Enum):
    """问题类型枚举"""
    PERFORMANCE_DEGRADATION = "performance_degradation"
    SECURITY_BREACH = "security_breach"
    NETWORK_CONNECTIVITY = "network_connectivity"
    MEMORY_LEAK = "memory_leak"
    DISK_SPACE_ISSUE = "disk_space_issue"
    SERVICE_UNAVAILABLE = "service_unavailable"
    DATABASE_ISSUE = "database_issue"
    APPLICATION_ERROR = "application_error"

class AgentType(Enum):
    """Agent类型枚举"""
    PERFORMANCE_AGENT = "performance_agent"
    SECURITY_AGENT = "security_agent"
    NETWORK_AGENT = "network_agent"
    MEMORY_AGENT = "memory_agent"
    DATABASE_AGENT = "database_agent"
    SYSTEM_AGENT = "system_agent"
    APPLICATION_AGENT = "application_agent"
    LOG_AGENT = "log_agent"

@dataclass
class LLMAnalysisResult:
    """LLM分析结果数据类"""
    request_id: str
    primary_issue: str
    root_cause: str
    severity_level: float
    confidence_score: float
    business_impact: str
    estimated_resolution_time: str
    anomaly_analysis: List[Dict[str, Any]]
    recommendations: Dict[str, List[Dict[str, Any]]]
    affected_services: Dict[str, Any]
    langraph_workflow: Dict[str, Any]

@dataclass
class AgentTask:
    """Agent任务数据类"""
    agent_type: AgentType
    task_description: str
    priority: str
    estimated_duration: str
    risk_level: str
    affected_services: List[str]
    parameters: Dict[str, Any]

class MCPLLMIntegrationManager:
    """MCP LLM集成管理器"""
    
    def __init__(self, mcp_base_url: str = "http://localhost:8085"):
        self.mcp_base_url = mcp_base_url
        self.session = None
        
        # Agent选择映射表
        self.agent_selection_map = {
            IssueType.PERFORMANCE_DEGRADATION: [
                AgentType.PERFORMANCE_AGENT,
                AgentType.SYSTEM_AGENT,
                AgentType.DATABASE_AGENT
            ],
            IssueType.SECURITY_BREACH: [
                AgentType.SECURITY_AGENT,
                AgentType.NETWORK_AGENT,
                AgentType.LOG_AGENT
            ],
            IssueType.NETWORK_CONNECTIVITY: [
                AgentType.NETWORK_AGENT,
                AgentType.SYSTEM_AGENT
            ],
            IssueType.MEMORY_LEAK: [
                AgentType.MEMORY_AGENT,
                AgentType.PERFORMANCE_AGENT,
                AgentType.APPLICATION_AGENT
            ],
            IssueType.DATABASE_ISSUE: [
                AgentType.DATABASE_AGENT,
                AgentType.PERFORMANCE_AGENT,
                AgentType.SYSTEM_AGENT
            ],
            IssueType.DISK_SPACE_ISSUE: [
                AgentType.SYSTEM_AGENT,
                AgentType.PERFORMANCE_AGENT
            ],
            IssueType.SERVICE_UNAVAILABLE: [
                AgentType.SYSTEM_AGENT,
                AgentType.NETWORK_AGENT,
                AgentType.APPLICATION_AGENT
            ],
            IssueType.APPLICATION_ERROR: [
                AgentType.APPLICATION_AGENT,
                AgentType.LOG_AGENT,
                AgentType.PERFORMANCE_AGENT
            ]
        }
    
    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()
    
    def parse_llm_result(self, llm_response: Dict[str, Any]) -> LLMAnalysisResult:
        """解析LLM分析结果"""
        try:
            analysis_result = llm_response["analysis_result"]
            overall_assessment = analysis_result["overall_assessment"]
            
            return LLMAnalysisResult(
                request_id=analysis_result["request_id"],
                primary_issue=overall_assessment["primary_issue"],
                root_cause=overall_assessment["root_cause"],
                severity_level=overall_assessment["severity_level"],
                confidence_score=overall_assessment["confidence_score"],
                business_impact=overall_assessment["business_impact"],
                estimated_resolution_time=overall_assessment["estimated_resolution_time"],
                anomaly_analysis=analysis_result["anomaly_analysis"],
                recommendations=analysis_result["recommendations"],
                affected_services=analysis_result["affected_services_summary"],
                langraph_workflow=analysis_result["langraph_workflow"]
            )
        except Exception as e:
            raise Exception(f"解析LLM结果失败: {e}")
    
    def select_agents(self, llm_result: LLMAnalysisResult) -> List[AgentType]:
        """智能选择子Agent"""
        selected_agents = []
        
        # 根据主要问题类型选择Agent
        for anomaly in llm_result.anomaly_analysis:
            issue_type = IssueType(anomaly["analysis"]["issue_type"])
            if issue_type in self.agent_selection_map:
                selected_agents.extend(self.agent_selection_map[issue_type])
        
        # 去重并保持顺序
        unique_agents = []
        for agent in selected_agents:
            if agent not in unique_agents:
                unique_agents.append(agent)
        
        # 根据严重程度调整Agent优先级
        if llm_result.severity_level >= 8.0:
            # 高严重程度，添加更多Agent
            if AgentType.SYSTEM_AGENT not in unique_agents:
                unique_agents.insert(0, AgentType.SYSTEM_AGENT)
        
        return unique_agents[:5]  # 最多选择5个Agent
    
    def distribute_tasks(self, selected_agents: List[AgentType], 
                        llm_result: LLMAnalysisResult) -> List[AgentTask]:
        """为选中的Agent分配任务"""
        tasks = []
        
        # 处理立即执行的操作
        for action in llm_result.recommendations.get("immediate_actions", []):
            # 根据操作类型选择对应的Agent
            agent_type = self._map_action_to_agent(action["action"])
            if agent_type in selected_agents:
                task = AgentTask(
                    agent_type=agent_type,
                    task_description=action["description"],
                    priority=action["priority"],
                    estimated_duration=action["estimated_duration"],
                    risk_level=action["risk_level"],
                    affected_services=action["affected_services"],
                    parameters={
                        "action": action["action"],
                        "severity": llm_result.severity_level,
                        "confidence": llm_result.confidence_score,
                        "root_cause": llm_result.root_cause
                    }
                )
                tasks.append(task)
        
        # 如果没有立即操作，为每个Agent创建分析任务
        if not tasks:
            for agent in selected_agents:
                task = AgentTask(
                    agent_type=agent,
                    task_description=f"分析{llm_result.primary_issue}问题",
                    priority="high" if llm_result.severity_level >= 7.0 else "medium",
                    estimated_duration="10 minutes",
                    risk_level="low",
                    affected_services=llm_result.affected_services.get("primary_affected", []),
                    parameters={
                        "primary_issue": llm_result.primary_issue,
                        "root_cause": llm_result.root_cause,
                        "severity": llm_result.severity_level,
                        "confidence": llm_result.confidence_score
                    }
                )
                tasks.append(task)
        
        return tasks
    
    def _map_action_to_agent(self, action: str) -> AgentType:
        """将操作映射到对应的Agent"""
        action_mapping = {
            "restart_database_connection_pool": AgentType.DATABASE_AGENT,
            "optimize_database_queries": AgentType.DATABASE_AGENT,
            "increase_connection_pool_size": AgentType.DATABASE_AGENT,
            "restart_application_service": AgentType.APPLICATION_AGENT,
            "check_memory_leak": AgentType.MEMORY_AGENT,
            "optimize_memory_usage": AgentType.MEMORY_AGENT,
            "check_network_connectivity": AgentType.NETWORK_AGENT,
            "optimize_network_routing": AgentType.NETWORK_AGENT,
            "check_security_vulnerabilities": AgentType.SECURITY_AGENT,
            "block_suspicious_ip": AgentType.SECURITY_AGENT,
            "analyze_system_performance": AgentType.PERFORMANCE_AGENT,
            "optimize_system_resources": AgentType.PERFORMANCE_AGENT,
            "check_disk_space": AgentType.SYSTEM_AGENT,
            "cleanup_old_files": AgentType.SYSTEM_AGENT,
            "analyze_logs": AgentType.LOG_AGENT
        }
        
        return action_mapping.get(action, AgentType.SYSTEM_AGENT)
    
    async def execute_agent_tasks(self, tasks: List[AgentTask]) -> Dict[str, Any]:
        """执行Agent任务"""
        results = {
            "plan_id": f"plan_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            "status": "executing",
            "tasks": [],
            "generated_scripts": [],
            "execution_time": 0.0
        }
        
        start_time = datetime.now()
        
        for i, task in enumerate(tasks):
            print(f"🤖 执行任务 {i+1}/{len(tasks)}: {task.agent_type.value}")
            
            # 调用MCP的auto_process_and_execute接口
            task_data = {
                "llm_data": {
                    "anomaly_type": task.agent_type.value.replace("_agent", ""),
                    "severity": task.parameters.get("severity", 5.0),
                    "description": task.task_description,
                    "affected_services": task.affected_services,
                    "suggested_actions": [task.task_description],
                    "confidence": task.parameters.get("confidence", 0.8),
                    "reasoning": f"由{task.agent_type.value}处理{task.parameters.get('primary_issue', 'unknown')}问题"
                }
            }
            
            try:
                url = f"{self.mcp_base_url}/api/intelligent/auto_process_and_execute"
                if self.session is None:
                    raise Exception("会话未初始化")
                
                async with self.session.post(url, json=task_data) as response:
                    if response.status == 200:
                        task_result = await response.json()
                        results["tasks"].append({
                            "task_id": f"task_{i+1}",
                            "agent_type": task.agent_type.value,
                            "status": task_result.get("status", "unknown"),
                            "result": task_result
                        })
                        
                        # 收集生成的脚本
                        if "generated_scripts" in task_result:
                            results["generated_scripts"].extend(task_result["generated_scripts"])
                        
                        print(f"✅ 任务完成: {task_result.get('plan_id', 'N/A')}")
                    else:
                        error_text = await response.text()
                        print(f"❌ 任务失败: {response.status} - {error_text}")
                        results["tasks"].append({
                            "task_id": f"task_{i+1}",
                            "agent_type": task.agent_type.value,
                            "status": "failed",
                            "error": f"HTTP {response.status}: {error_text}"
                        })
            
            except Exception as e:
                print(f"❌ 任务执行异常: {e}")
                results["tasks"].append({
                    "task_id": f"task_{i+1}",
                    "agent_type": task.agent_type.value,
                    "status": "failed",
                    "error": str(e)
                })
        
        # 计算执行时间
        end_time = datetime.now()
        results["execution_time"] = (end_time - start_time).total_seconds()
        
        # 判断整体状态
        success_count = sum(1 for task in results["tasks"] if task["status"] == "completed")
        if success_count == len(tasks):
            results["status"] = "completed"
        elif success_count > 0:
            results["status"] = "partial"
        else:
            results["status"] = "failed"
        
        return results
    
    async def deploy_to_kylin(self, ansible_scripts: List[str]) -> Dict[str, Any]:
        """部署Ansible脚本到Kylin"""
        if not ansible_scripts:
            return {"status": "no_scripts", "message": "没有可部署的脚本"}
        
        deployment_results = []
        
        for i, script in enumerate(ansible_scripts):
            try:
                # 调用MCP的deploy_to_kylin接口
                deploy_data = {
                    "script_content": script,
                    "script_name": f"deployment_script_{i+1}.yml",
                    "target_environment": "production"
                }
                
                url = f"{self.mcp_base_url}/api/intelligent/deploy_to_kylin"
                if self.session is None:
                    raise Exception("会话未初始化")
                
                async with self.session.post(url, json=deploy_data) as response:
                    if response.status == 200:
                        deploy_result = await response.json()
                        deployment_results.append({
                            "script_id": i+1,
                            "status": "deployed",
                            "result": deploy_result
                        })
                        print(f"✅ 脚本 {i+1} 部署成功")
                    else:
                        error_text = await response.text()
                        deployment_results.append({
                            "script_id": i+1,
                            "status": "failed",
                            "error": f"HTTP {response.status}: {error_text}"
                        })
                        print(f"❌ 脚本 {i+1} 部署失败: {error_text}")
            
            except Exception as e:
                deployment_results.append({
                    "script_id": i+1,
                    "status": "failed",
                    "error": str(e)
                })
                print(f"❌ 脚本 {i+1} 部署异常: {e}")
        
        return {
            "status": "completed",
            "deployment_results": deployment_results,
            "total_scripts": len(ansible_scripts),
            "successful_deployments": sum(1 for r in deployment_results if r["status"] == "deployed")
        }
    
    async def process_llm_analysis_result(self, llm_response: Dict[str, Any]) -> Dict[str, Any]:
        """处理LLM分析结果的完整流程"""
        try:
            print("🔄 开始处理LLM分析结果...")
            
            # 1. 解析LLM结果
            print("📋 解析LLM分析结果...")
            llm_result = self.parse_llm_result(llm_response)
            print(f"✅ 解析完成: {llm_result.primary_issue} (严重程度: {llm_result.severity_level})")
            
            # 2. 智能选择子Agent
            print("🤖 智能选择子Agent...")
            selected_agents = self.select_agents(llm_result)
            print(f"✅ 选中Agent: {[agent.value for agent in selected_agents]}")
            
            # 3. 分配任务
            print("📝 分配任务...")
            tasks = self.distribute_tasks(selected_agents, llm_result)
            print(f"✅ 分配任务: {len(tasks)} 个任务")
            
            # 4. 执行任务
            print("⚡ 执行Agent任务...")
            execution_result = await self.execute_agent_tasks(tasks)
            print(f"✅ 任务执行完成: {execution_result['status']}")
            
            # 5. 部署到Kylin
            if execution_result["generated_scripts"]:
                print("🚀 部署到Kylin...")
                deployment_result = await self.deploy_to_kylin(execution_result["generated_scripts"])
                print(f"✅ 部署完成: {deployment_result['successful_deployments']}/{deployment_result['total_scripts']} 成功")
            else:
                deployment_result = {"status": "no_scripts"}
            
            return {
                "success": True,
                "llm_result": llm_result,
                "selected_agents": [agent.value for agent in selected_agents],
                "execution_result": execution_result,
                "deployment_result": deployment_result,
                "workflow_status": "completed"
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "workflow_status": "failed"
            }

# 使用示例
async def example_usage():
    """使用示例"""
    # 模拟LLM+RAG+Langraph的分析结果
    llm_response = {
        "success": True,
        "analysis_result": {
            "request_id": "req_1703123456",
            "analysis_timestamp": "2023-12-21T10:32:30.000000+00:00",
            "overall_assessment": {
                "primary_issue": "database_connection_pool_exhaustion",
                "root_cause": "High CPU usage caused by inefficient database queries",
                "severity_level": 7.5,
                "confidence_score": 0.92,
                "business_impact": "moderate",
                "estimated_resolution_time": "30 minutes"
            },
            "anomaly_analysis": [
                {
                    "anomaly_id": "cpu_usage_system_server-01_1703123456",
                    "analysis": {
                        "issue_type": "performance_degradation",
                        "root_cause": "Database connection pool exhaustion due to slow queries",
                        "severity": 7.5,
                        "confidence": 0.92
                    }
                }
            ],
            "recommendations": {
                "immediate_actions": [
                    {
                        "action": "restart_database_connection_pool",
                        "description": "重启数据库连接池以释放资源",
                        "priority": "high",
                        "estimated_duration": "5 minutes",
                        "risk_level": "low",
                        "affected_services": ["database"]
                    },
                    {
                        "action": "optimize_database_queries",
                        "description": "优化慢查询，添加必要的索引",
                        "priority": "high",
                        "estimated_duration": "20 minutes",
                        "risk_level": "medium",
                        "affected_services": ["database", "api"]
                    }
                ]
            },
            "affected_services_summary": {
                "primary_affected": ["database", "api"],
                "secondary_affected": ["web"]
            },
            "langraph_workflow": {
                "nodes_executed": ["anomaly_detection", "root_cause_analysis"],
                "reasoning_chain": ["Detected CPU usage spike", "Identified connection pool exhaustion"]
            }
        }
    }
    
    # 创建MCP集成管理器
    async with MCPLLMIntegrationManager() as manager:
        # 处理LLM分析结果
        result = await manager.process_llm_analysis_result(llm_response)
        
        if result["success"]:
            print("🎉 LLM+MCP集成工作流执行成功!")
            print(f"选中的Agent: {result['selected_agents']}")
            print(f"执行状态: {result['execution_result']['status']}")
            print(f"部署结果: {result['deployment_result']['status']}")
        else:
            print(f"❌ 集成工作流执行失败: {result['error']}")

if __name__ == "__main__":
    asyncio.run(example_usage()) 